import random

from gxl_ai_utils.utils import utils_file
from tqdm import tqdm


def get_token_text_jsonl():
    input_token_scp = "/home/work_nfs8/xlgeng/new_workspace/wenet_gxl_salmonn_TTS/examples/aishell/wenetspeech4tts_handler/data_list/train_token.scp"
    input_text_scp = "/home/work_nfs8/xlgeng/new_workspace/wenet_gxl_salmonn_TTS/examples/aishell/wenetspeech4tts_handler/data_list/train_text.scp"
    token_dict = utils_file.load_dict_from_scp(input_token_scp)
    text_dict = utils_file.load_dict_from_scp(input_text_scp)
    res_dict_list = []
    for key in tqdm(token_dict.keys(), desc="token_dict", total=len(token_dict)):
        temp_dict = {}
        if key not in text_dict.keys():
            utils_file.logging_print("warning: key:{} not in text_dict".format(key))
            continue
        temp_dict["key"] = key
        temp_dict["token"] = token_dict[key]
        temp_dict["txt"] = text_dict[key]
        res_dict_list.append(temp_dict)

    output_jsonl = "/home/work_nfs8/xlgeng/new_workspace/wenet_gxl_salmonn_TTS/examples/aishell/wenetspeech4tts_handler/data_list/train_token_text.jsonl"
    utils_file.write_dict_list_to_jsonl(res_dict_list, output_jsonl)

def split_train_cv_text_jsonl():
    input_jsonl_path = "/home/work_nfs8/xlgeng/new_workspace/wenet_gxl_salmonn_TTS/examples/aishell/wenetspeech4tts_handler/data_list/train_token_text.jsonl"
    cv_path = "/home/work_nfs8/xlgeng/new_workspace/wenet_gxl_salmonn_TTS/examples/aishell/wenetspeech4tts_handler/data_list/cv_token_text.jsonl"
    test_path = "/home/work_nfs8/xlgeng/new_workspace/wenet_gxl_salmonn_TTS/examples/aishell/wenetspeech4tts_handler/data_list/test_token_text.jsonl"
    all_lines = utils_file.load_list_file_clean(input_jsonl_path)
    random.shuffle(all_lines)
    train_lines, cv_lins, test_lines = all_lines[1500:], all_lines[:1000], all_lines[1000:1500]
    utils_file.write_list_to_file(train_lines, input_jsonl_path)
    utils_file.write_list_to_file(cv_lins, cv_path)
    utils_file.write_list_to_file(test_lines, test_path)


if __name__ == "__main__":
    split_train_cv_text_jsonl()
